use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{Mutex, RwLock};
use tokio::sync::Notify;
use std::sync::Arc;

type SyncPoint = u64;

pub struct SyncManager {
    next_id: AtomicU64,
    id_map: RwLock<HashMap<u64, Arc<Notify>>>,
    response_map: RwLock<HashMap<u64, u64>>,
}

impl SyncManager {
    pub fn new() -> Self {
        SyncManager {
            next_id: AtomicU64::new(1),
            id_map: RwLock::new(HashMap::new()),
            response_map: RwLock::new(HashMap::new()),
        }
    }

    pub async fn generate_unique_sync_point(&self) -> SyncPoint {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);
        let mut id_map = self.id_map.write().await;
        id_map.insert(id, Arc::new(Notify::new()));
        id
    }

    pub async fn wait_for_sync_point(&self, sync_point_id: SyncPoint) -> u64 {
        let notify = {
            let id_map = self.id_map.read().await;
            id_map.get(&sync_point_id).cloned()
        };

        if let Some(notify) = notify {
            notify.notified().await;
        }

        let mut response_map = self.response_map.write().await;
        response_map.remove(&sync_point_id).unwrap_or(0)
    }

    pub async fn release_sync_point(&self, sync_point_id: SyncPoint, id: u64) {
        let mut id_map = self.id_map.write().await;
        if let Some(notify) = id_map.remove(&sync_point_id) {
            let mut response_map = self.response_map.write().await;
            response_map.insert(sync_point_id, id);
            notify.notify_waiters();
        }
    }

    pub async fn stop(&self) {
        let keys: Vec<u64> = {
            let id_map = self.id_map.read().await;
            id_map.keys().cloned().collect()
        };

        for key in keys {
            self.release_sync_point(key, 0).await;
        }
    }
}

impl Drop for SyncManager {
    fn drop(&mut self) {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            self.stop().await;
        });
    }
}